{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "777fc40d",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\" width=\"90px\">\n",
    "\n",
    "# PySpark Huggingface Inferencing\n",
    "### Conditional generation with Tensorflow\n",
    "\n",
    "In this notebook, we demonstrate distributed inference with the T5 transformer to perform sentence translation.  \n",
    "From: https://huggingface.co/docs/transformers/model_doc/t5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05c79ac4-bf25-421e-b55e-020d6d9e15d5",
   "metadata": {},
   "source": [
    "Note that cuFFT/cuDNN/cuBLAS registration errors are expected (as of `tf=2.17.0`) and will not affect behavior, as noted in [this issue.](https://github.com/tensorflow/tensorflow/issues/62075)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f6f0dbf3-712b-4c58-85eb-261ce15bb2be",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:53:50.831324: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-02-04 13:53:50.838528: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2025-02-04 13:53:50.846226: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2025-02-04 13:53:50.848585: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2025-02-04 13:53:50.854859: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2025-02-04 13:53:51.229622: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, TFT5ForConditionalGeneration\n",
    "\n",
    "# Manually enable Huggingface tokenizer parallelism to avoid disabling with PySpark parallelism.\n",
    "# See (https://github.com/huggingface/transformers/issues/5486) for more info. \n",
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "275890d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.17.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1738706031.770264 3625306 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706031.793270 3625306 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706031.796251 3625306 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "# Enable GPU memory growth\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "if gpus:\n",
    "    try:\n",
    "        for gpu in gpus:\n",
    "            tf.config.experimental.set_memory_growth(gpu, True)\n",
    "    except RuntimeError as e:\n",
    "        print(e)\n",
    "        \n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2684fb41-9467-40c0-9d7e-a1cc867c5a3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0000 00:00:1738706032.132191 3625306 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706032.134996 3625306 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706032.137528 3625306 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706032.251302 3625306 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706032.252345 3625306 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706032.253281 3625306 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2025-02-04 13:53:52.254192: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 43462 MB memory:  -> device: 0, name: NVIDIA RTX A6000, pci bus id: 0000:01:00.0, compute capability: 8.6\n",
      "All PyTorch model weights were used when initializing TFT5ForConditionalGeneration.\n",
      "\n",
      "All the weights of TFT5ForConditionalGeneration were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFT5ForConditionalGeneration for predictions without further training.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"google-t5/t5-small\")\n",
    "model = TFT5ForConditionalGeneration.from_pretrained(\"google-t5/t5-small\")\n",
    "\n",
    "task_prefix = \"translate English to German: \"\n",
    "\n",
    "lines = [\n",
    "    \"The house is wonderful\",\n",
    "    \"Welcome to NYC\",\n",
    "    \"HuggingFace is a company\"\n",
    "]\n",
    "\n",
    "input_sequences = [task_prefix + l for l in lines]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6eb2dfdb-0ad3-4d0f-81a4-268d92c53759",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1738706033.555987 3625654 service.cc:146] XLA service 0x712d300025f0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:\n",
      "I0000 00:00:1738706033.556005 3625654 service.cc:154]   StreamExecutor device (0): NVIDIA RTX A6000, Compute Capability 8.6\n",
      "2025-02-04 13:53:53.558887: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:268] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.\n",
      "2025-02-04 13:53:53.569767: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:531] Loaded cuDNN version 8907\n",
      "I0000 00:00:1738706033.604327 3625654 device_compiler.h:188] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.\n"
     ]
    }
   ],
   "source": [
    "inputs = tokenizer(input_sequences, \n",
    "                      padding=True,\n",
    "                      return_tensors=\"tf\")\n",
    "outputs = model.generate(input_ids=inputs[\"input_ids\"], attention_mask=inputs[\"attention_mask\"], max_length=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "720158d4-e0e0-4904-b096-e5aede756afd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Das Haus ist wunderbar',\n",
       " 'Willkommen in NYC',\n",
       " 'HuggingFace ist ein Unternehmen']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[tokenizer.decode(o, skip_special_tokens=True) for o in outputs]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "546eabe0",
   "metadata": {},
   "source": [
    "## PySpark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "68121304-f1df-466e-9347-c9d2b36a9b3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.types import *\n",
    "from pyspark import SparkConf\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.functions import pandas_udf, col, struct\n",
    "from pyspark.ml.functions import predict_batch_udf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2f6db1f0-7d68-4af7-8bd6-c9fa45906c61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import datasets\n",
    "from datasets import load_dataset\n",
    "datasets.disable_progress_bars()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d636975",
   "metadata": {},
   "source": [
    "Check the cluster environment to handle any platform-specific Spark configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ca351245",
   "metadata": {},
   "outputs": [],
   "source": [
    "on_databricks = os.environ.get(\"DATABRICKS_RUNTIME_VERSION\", False)\n",
    "on_dataproc = os.environ.get(\"DATAPROC_IMAGE_VERSION\", False)\n",
    "on_standalone = not (on_databricks or on_dataproc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3199f8b",
   "metadata": {},
   "source": [
    "#### Create Spark Session\n",
    "\n",
    "For local standalone clusters, we'll connect to the cluster and create the Spark Session.  \n",
    "For CSP environments, Spark will either be preconfigured (Databricks) or we'll need to create the Spark Session (Dataproc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6279a849",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:53:54 WARN Utils: Your hostname, cb4ae00-lcedt resolves to a loopback address: 127.0.1.1; using 10.110.47.100 instead (on interface eno1)\n",
      "25/02/04 13:53:54 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
      "25/02/04 13:53:55 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
     ]
    }
   ],
   "source": [
    "conf = SparkConf()\n",
    "\n",
    "if 'spark' not in globals():\n",
    "    if on_standalone:\n",
    "        import socket\n",
    "        \n",
    "        conda_env = os.environ.get(\"CONDA_PREFIX\")\n",
    "        hostname = socket.gethostname()\n",
    "        conf.setMaster(f\"spark://{hostname}:7077\")\n",
    "        conf.set(\"spark.pyspark.python\", f\"{conda_env}/bin/python\")\n",
    "        conf.set(\"spark.pyspark.driver.python\", f\"{conda_env}/bin/python\")\n",
    "    elif on_dataproc:\n",
    "        conf.set(\"spark.executorEnv.TF_GPU_ALLOCATOR\", \"cuda_malloc_async\")\n",
    "\n",
    "    conf.set(\"spark.executor.cores\", \"8\")\n",
    "    conf.set(\"spark.task.resource.gpu.amount\", \"0.125\")\n",
    "    conf.set(\"spark.executor.resource.gpu.amount\", \"1\")\n",
    "    conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", \"true\")\n",
    "    conf.set(\"spark.python.worker.reuse\", \"true\")\n",
    "\n",
    "conf.set(\"spark.sql.execution.arrow.maxRecordsPerBatch\", \"1000\")\n",
    "spark = SparkSession.builder.appName(\"spark-dl-examples\").config(conf=conf).getOrCreate()\n",
    "sc = spark.sparkContext"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f311650",
   "metadata": {},
   "source": [
    "Load the IMBD Movie Reviews dataset from Huggingface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b8453111-d068-49bb-ab91-8ae3d8bcdb7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"imdb\", split=\"test\")\n",
    "dataset = dataset.to_pandas().drop(columns=\"label\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fd5b472-47e8-4804-9907-772793fedb2b",
   "metadata": {},
   "source": [
    "### Create PySpark DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d24d9404-0269-476e-a9dd-1842667c915a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StructType([StructField('text', StringType(), True)])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = spark.createDataFrame(dataset).repartition(8)\n",
    "df.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c76314b7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "25000"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4384c762-1f79-4f60-876c-94b1f552e8fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:54:01 WARN TaskSetManager: Stage 6 contains a task of very large size (4021 KiB). The maximum recommended task size is 1000 KiB.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[Row(text=\"Anyone remember the first CKY, CKY2K etc..? Back when it was about making crazy cool stuff, rather than watching Bam Margera act like a douchebag, spoiled 5 year old, super/rock-star wannabe.<br /><br />The show used to be awesome, however, Bam's fame and wealth has led him to believe, that we now enjoy him acting childish and idiotic, more than actual cool stuff, that used to be in ex. CKY2K.<br /><br />The acts are so repetitive, there's like nothing new, except annoying stupidity and rehearsed comments... The only things we see is Bam Margera, so busy showing us how much he doesn't care, how much money he got or whatsoever.<br /><br />I really got nothing much left to say except, give us back CKY2K, cause Bam suck..<br /><br />I enjoy watching Steve-o, Knoxville etc. a thousand times more.\")]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.take(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42ba3513-82dd-47e7-8193-eb4389458757",
   "metadata": {},
   "source": [
    "### Save the test dataset as parquet files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e7eec8ec-4126-4890-b957-025809fad67d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:54:02 WARN TaskSetManager: Stage 9 contains a task of very large size (4021 KiB). The maximum recommended task size is 1000 KiB.\n"
     ]
    }
   ],
   "source": [
    "data_path = \"spark-dl-datasets/imdb_test\"\n",
    "if on_databricks:\n",
    "    dbutils.fs.mkdirs(\"/FileStore/spark-dl-datasets\")\n",
    "    data_path = \"dbfs:/FileStore/\" + data_path\n",
    "\n",
    "df.write.mode(\"overwrite\").parquet(data_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "078425e1",
   "metadata": {},
   "source": [
    "#### Load and preprocess DataFrame\n",
    "\n",
    "Define our preprocess function. We'll take the first sentence from each sample as our input for translation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b9a0889a-35b4-493a-8197-1146fc7efd53",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(text: pd.Series, prefix: str = \"\") -> pd.Series:\n",
    "    @pandas_udf(\"string\")\n",
    "    def _preprocess(text: pd.Series) -> pd.Series:\n",
    "        return pd.Series([prefix + s.split(\".\")[0] for s in text])\n",
    "    return _preprocess(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c483e4d4-9ab1-416f-a766-694e17490fd3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------------------------------------------+\n",
      "|                                                                                                text|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "|Doesn't anyone bother to check where this kind of sludge comes from before blathering on about it...|\n",
      "|There were two things I hated about WASTED : The directing and the script . I know I`m opening my...|\n",
      "|I'm rather surprised that anybody found this film touching or moving.<br /><br />The basic premis...|\n",
      "|Cultural Vandalism Is the new Hallmark production of Gulliver's Travels an act of cultural vandal...|\n",
      "|I was at Wrestlemania VI in Toronto as a 10 year old, and the event I saw then was pretty differe...|\n",
      "|This movie has been done before. It is basically a unoriginal combo of \"Napoleon Dynamite\" and \"S...|\n",
      "|[ as a new resolution for this year 2005, i decide to write a comment for each movie I saw in the...|\n",
      "|This movie is over hyped!! I am sad to say that I manage to watch the first 15 minutes of this mo...|\n",
      "|This show had a promising start as sort of the opposite of 'Oceans 11' but has developed into a s...|\n",
      "|MINOR PLOT SPOILERS AHEAD!!!<br /><br />How did such talented actors get involved in such mindles...|\n",
      "|There is not one character on this sitcom with any redeeming qualities. They are all self-centere...|\n",
      "|Tommy Lee Jones was the best Woodroe and no one can play Woodroe F. Call better than he. Not only...|\n",
      "|My wife rented this movie and then conveniently never got to see it. If I ever want to torture he...|\n",
      "|This is one of those star-filled over-the-top comedies that could a) be hysterical, or b) wish th...|\n",
      "|This excruciatingly boring and unfunny movie made me think that Chaplin was the real Hitler, as o...|\n",
      "|you will likely be sorely disappointed by this sequel that's not a sequel.AWIL is a classic.but t...|\n",
      "|If I was British, I would be embarrassed by this portrayal of incompetence. A protection agent of...|\n",
      "|One of those movies in which there are no big twists whatsoever and you can predict pretty much w...|\n",
      "|This show is like watching someone who is in training to someday host a show. There are some good...|\n",
      "|Sigh. I'm baffled when I see a short like this get attention and assignments and whatnot. I saw t...|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Limit to N rows, since this can be slow\n",
    "df = spark.read.parquet(data_path).limit(256).repartition(8)\n",
    "df.show(truncate=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9f8e538",
   "metadata": {},
   "source": [
    "Append a prefix to tell the model to translate English to French:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "831bc52c-a5c6-4c29-a6da-0566b5167773",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------------------------------------------+\n",
      "|                                                                                               input|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "|translate English to French: Doesn't anyone bother to check where this kind of sludge comes from ...|\n",
      "|translate English to French: There were two things I hated about WASTED : The directing and the s...|\n",
      "|   translate English to French: I'm rather surprised that anybody found this film touching or moving|\n",
      "|translate English to French: Cultural Vandalism Is the new Hallmark production of Gulliver's Trav...|\n",
      "|translate English to French: I was at Wrestlemania VI in Toronto as a 10 year old, and the event ...|\n",
      "|                                        translate English to French: This movie has been done before|\n",
      "|translate English to French: [ as a new resolution for this year 2005, i decide to write a commen...|\n",
      "|translate English to French: This movie is over hyped!! I am sad to say that I manage to watch th...|\n",
      "|translate English to French: This show had a promising start as sort of the opposite of 'Oceans 1...|\n",
      "|translate English to French: MINOR PLOT SPOILERS AHEAD!!!<br /><br />How did such talented actors...|\n",
      "| translate English to French: There is not one character on this sitcom with any redeeming qualities|\n",
      "|     translate English to French: Tommy Lee Jones was the best Woodroe and no one can play Woodroe F|\n",
      "|    translate English to French: My wife rented this movie and then conveniently never got to see it|\n",
      "|translate English to French: This is one of those star-filled over-the-top comedies that could a)...|\n",
      "|translate English to French: This excruciatingly boring and unfunny movie made me think that Chap...|\n",
      "|translate English to French: you will likely be sorely disappointed by this sequel that's not a s...|\n",
      "|translate English to French: If I was British, I would be embarrassed by this portrayal of incomp...|\n",
      "|translate English to French: One of those movies in which there are no big twists whatsoever and ...|\n",
      "|translate English to French: This show is like watching someone who is in training to someday hos...|\n",
      "|                                                                   translate English to French: Sigh|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "input_df = df.select(preprocess(col(\"text\"), \"translate English to French: \").alias(\"input\")).cache()\n",
    "input_df.show(truncate=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec53a65c",
   "metadata": {},
   "source": [
    "## Inference using Spark DL API\n",
    "\n",
    "Distributed inference using the PySpark [predict_batch_udf](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.functions.predict_batch_udf.html#pyspark.ml.functions.predict_batch_udf):\n",
    "\n",
    "- predict_batch_fn uses Tensorflow APIs to load the model and return a predict function which operates on numpy arrays \n",
    "- predict_batch_udf will convert the Spark DataFrame columns into numpy input batches for the predict function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e7ae69d3-70c2-4765-928f-c96a7ba59829",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_batch_fn():\n",
    "    import tensorflow as tf\n",
    "    import numpy as np\n",
    "    from transformers import TFT5ForConditionalGeneration, AutoTokenizer\n",
    "\n",
    "    # Enable GPU memory growth\n",
    "    print(\"initializing model\")\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        except RuntimeError as e:\n",
    "            print(e)\n",
    "\n",
    "    model = TFT5ForConditionalGeneration.from_pretrained(\"google-t5/t5-small\")\n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"google-t5/t5-small\")\n",
    "\n",
    "    def predict(inputs):\n",
    "        flattened = np.squeeze(inputs).tolist()\n",
    "        inputs = tokenizer(flattened, \n",
    "                            padding=True, \n",
    "                            return_tensors=\"tf\")\n",
    "        outputs = model.generate(input_ids=inputs[\"input_ids\"],\n",
    "                                 attention_mask=inputs[\"attention_mask\"],\n",
    "                                 max_length=128)\n",
    "        string_outputs = np.array([tokenizer.decode(o, skip_special_tokens=True) for o in outputs])\n",
    "        print(\"predict: {}\".format(len(flattened)))\n",
    "        return string_outputs\n",
    "    \n",
    "    return predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "36684f59-d947-43f8-a2e8-c7a423764e88",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate = predict_batch_udf(predict_batch_fn,\n",
    "                             return_type=StringType(),\n",
    "                             batch_size=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6a01c855-8fa1-4765-a3a5-2c9dd872df10",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 24:====================================>                     (5 + 3) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 9.07 ms, sys: 8.83 ms, total: 17.9 ms\n",
      "Wall time: 19.3 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "preds = input_df.withColumn(\"preds\", generate(struct(\"input\")))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d912d4b0-cd0b-44ea-859a-b23455cc2700",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 27:==================================================>       (7 + 1) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.51 ms, sys: 4.96 ms, total: 12.5 ms\n",
      "Wall time: 12.4 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df.withColumn(\"preds\", generate(\"input\"))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "5fe3d88b-30f7-468f-8db8-1f4118d0f26c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 30:=====================>                                    (3 + 5) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.46 ms, sys: 5.98 ms, total: 11.4 ms\n",
      "Wall time: 11.4 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df.withColumn(\"preds\", generate(col(\"input\")))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4ad9b365-4b9a-438e-8fdf-47da55cb1cf4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 33:>                                                         (0 + 1) / 1]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|                                             input|                                             preds|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|translate English to French: Doesn't anyone bot...|Ne s'ennuie-t-il pas de vérifier où viennent ce...|\n",
      "|translate English to French: There were two thi...|Il y avait deux choses que j'ai hâte de voir : ...|\n",
      "|translate English to French: I'm rather surpris...|Je suis plutôt surpris que quelqu'un ait trouvé...|\n",
      "|translate English to French: Cultural Vandalism...|Vandalisme culturel La nouvelle production Hall...|\n",
      "|translate English to French: I was at Wrestlema...|J'étais à Wrestlemania VI à Toronto en 10 ans, ...|\n",
      "|translate English to French: This movie has bee...|                       Ce film a été réalisé avant|\n",
      "|translate English to French: [ as a new resolut...|[ en tant que nouvelle résolution pour cette an...|\n",
      "|translate English to French: This movie is over...|Je suis triste de dire que je parviens à regard...|\n",
      "|translate English to French: This show had a pr...|Ce spectacle a eu un début prometteur en l'espè...|\n",
      "|translate English to French: MINOR PLOT SPOILER...|br />br /> Comment ces acteurs talentueux ont-i...|\n",
      "|translate English to French: There is not one c...|Il n'y a pas d'un personnage sur ce sitcom ayan...|\n",
      "|translate English to French: Tommy Lee Jones wa...|Tommy Lee Jones était le meilleur Woodroe et pe...|\n",
      "|translate English to French: My wife rented thi...|Ma femme a loué ce film et n'a jamais pu le voi...|\n",
      "|translate English to French: This is one of tho...|C’est l’une des comédies en étoiles à l’étoile ...|\n",
      "|translate English to French: This excruciatingl...|Ce film excruciant ennuyant et infaillible m’a ...|\n",
      "|translate English to French: you will likely be...|Vous serez probablement très déçu par cette séq...|\n",
      "|translate English to French: If I was British, ...|Si j'étais britannique, je seraitis embarrassé ...|\n",
      "|translate English to French: One of those movie...|Un des films dans lesquels il n'y a pas de gros...|\n",
      "|translate English to French: This show is like ...|Ce spectacle ressemble à l'observation d'une pe...|\n",
      "|                 translate English to French: Sigh|                                             Pesée|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "preds.show(truncate=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "1eb0c83b-d91b-4f8c-a5e7-c35f55c88108",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_df2 = df.select(preprocess(col(\"text\"), \"translate English to German: \").alias(\"input\")).cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "6f6b70f9-188a-402b-9143-78a5788140e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 36:==================================================>       (7 + 1) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 9.1 ms, sys: 4.04 ms, total: 13.1 ms\n",
      "Wall time: 14.9 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "preds = input_df2.withColumn(\"preds\", generate(struct(\"input\")))\n",
    "result = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "031a6a5e-7999-4653-b394-19ed478d8c96",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 39:==================================================>       (7 + 1) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.62 ms, sys: 5.23 ms, total: 11.9 ms\n",
      "Wall time: 11.9 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df2.withColumn(\"preds\", generate(\"input\"))\n",
    "result = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "229b6515-82f6-4e9c-90f0-a9c3cfb26301",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 42:==============>                                           (2 + 6) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8.67 ms, sys: 3.27 ms, total: 11.9 ms\n",
      "Wall time: 11.7 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df2.withColumn(\"preds\", generate(col(\"input\")))\n",
    "result = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8be750ac-fa39-452e-bb4c-c2270bc2f70d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 45:>                                                         (0 + 1) / 1]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|                                             input|                                             preds|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|translate English to German: Doesn't anyone bot...|Warum hat man sich nicht angeschaut, woher der ...|\n",
      "|translate English to German: There were two thi...|Es gab zwei Dinge, die ich hat an WASTED gehass...|\n",
      "|translate English to German: I'm rather surpris...|Ich bin ziemlich überrascht, dass jemand diesen...|\n",
      "|translate English to German: Cultural Vandalism...|Kultureller Vandalismus Ist die neue Hallmark-P...|\n",
      "|translate English to German: I was at Wrestlema...|Ich war als 10 Jahre alt bei Wrestlemania VI in...|\n",
      "|translate English to German: This movie has bee...|             Dieser Film wurde bereits vorgenommen|\n",
      "|translate English to German: [ as a new resolut...|[ als neue Entschließung für dieses Jahr 2005, ...|\n",
      "|translate English to German: This movie is over...|Ich hoffe, dass ich die ersten 15 Minuten diese...|\n",
      "|translate English to German: This show had a pr...|Diese Show hatte einen vielversprechenden Start...|\n",
      "|translate English to German: MINOR PLOT SPOILER...|br />br />Wie haben sich so talentierte Schausp...|\n",
      "|translate English to German: There is not one c...|Es gibt keinen Charakter auf dieser Seite mit i...|\n",
      "|translate English to German: Tommy Lee Jones wa...|Tommy Lee Jones war der beste Woodroe und niema...|\n",
      "|translate English to German: My wife rented thi...|Meine Frau hat diesen Film vermietet und dann b...|\n",
      "|translate English to German: This is one of tho...|Dies ist eines der Sterne-gefüllten über-the-to...|\n",
      "|translate English to German: This excruciatingl...|Dieser schreckliche langweilige und unfunnelnde...|\n",
      "|translate English to German: you will likely be...|Sie werden wahrscheinlich ernsthaft enttäuscht ...|\n",
      "|translate English to German: If I was British, ...|Wenn ich Britisch wäre, wäre ich beschämt über ...|\n",
      "|translate English to German: One of those movie...|Einer der Filme, in denen es keine großen Drehu...|\n",
      "|translate English to German: This show is like ...|Diese Show ist wie ein jemanden, der in Ausbild...|\n",
      "|                 translate English to German: Sigh|                                            Segnen|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "preds.show(truncate=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5803188",
   "metadata": {},
   "source": [
    "## Using Triton Inference Server\n",
    "In this section, we demonstrate integration with the [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server), an open-source, GPU-accelerated serving solution for DL.  \n",
    "We use [PyTriton](https://github.com/triton-inference-server/pytriton), a Flask-like framework that handles client/server communication with the Triton server.  \n",
    "\n",
    "The process looks like this:\n",
    "- Distribute a PyTriton task across the Spark cluster, instructing each node to launch a Triton server process.\n",
    "- Define a Triton inference function, which contains a client that binds to the local server on a given node and sends inference requests.\n",
    "- Wrap the Triton inference function in a predict_batch_udf to launch parallel inference requests using Spark.\n",
    "- Finally, distribute a shutdown signal to terminate the Triton server processes on each node.\n",
    "\n",
    "<img src=\"../images/spark-server.png\" alt=\"drawing\" width=\"700\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "6d09f972",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2964ffee",
   "metadata": {},
   "source": [
    "Import the helper class from server_utils.py:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "f1083dc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.addPyFile(\"server_utils.py\")\n",
    "\n",
    "from server_utils import TritonServerManager"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "066c8695",
   "metadata": {},
   "source": [
    "Define the Triton Server function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "afd00b7e-8150-4c95-a2e4-037e9c90f92a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_server(ports):\n",
    "    import time\n",
    "    import signal\n",
    "    import numpy as np\n",
    "    import tensorflow as tf\n",
    "    from transformers import TFT5ForConditionalGeneration, AutoTokenizer\n",
    "    from pytriton.decorators import batch\n",
    "    from pytriton.model_config import DynamicBatcher, ModelConfig, Tensor\n",
    "    from pytriton.triton import Triton, TritonConfig\n",
    "    from pyspark import TaskContext\n",
    "\n",
    "    print(f\"SERVER: Initializing Conditional Generation model on worker {TaskContext.get().partitionId()}.\")\n",
    "\n",
    "    # Enable GPU memory growth\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        except RuntimeError as e:\n",
    "            print(e)\n",
    "    \n",
    "    tokenizer = AutoTokenizer.from_pretrained(\"google-t5/t5-small\")\n",
    "    model = TFT5ForConditionalGeneration.from_pretrained(\"google-t5/t5-small\")\n",
    "\n",
    "    @batch\n",
    "    def _infer_fn(**inputs):\n",
    "        sentences = np.squeeze(inputs[\"text\"]).tolist()\n",
    "        print(f\"SERVER: Received batch of size {len(sentences)}\")\n",
    "        decoded_sentences = [s.decode(\"utf-8\") for s in sentences]\n",
    "        inputs = tokenizer(decoded_sentences,\n",
    "                            padding=True,\n",
    "                            return_tensors=\"tf\")\n",
    "        output_ids = model.generate(input_ids=inputs[\"input_ids\"],\n",
    "                                    attention_mask=inputs[\"attention_mask\"],\n",
    "                                    max_length=128)\n",
    "        outputs = np.array([[tokenizer.decode(o, skip_special_tokens=True)] for o in output_ids])\n",
    "        return {\n",
    "            \"translations\": outputs,\n",
    "        }\n",
    "\n",
    "    workspace_path = f\"/tmp/triton_{time.strftime('%m_%d_%M_%S')}\"\n",
    "    triton_conf = TritonConfig(http_port=ports[0], grpc_port=ports[1], metrics_port=ports[2])\n",
    "    with Triton(config=triton_conf, workspace=workspace_path) as triton:\n",
    "        triton.bind(\n",
    "            model_name=\"ConditionalGeneration\",\n",
    "            infer_func=_infer_fn,\n",
    "            inputs=[\n",
    "                Tensor(name=\"text\", dtype=object, shape=(-1,)),\n",
    "            ],\n",
    "            outputs=[\n",
    "                Tensor(name=\"translations\", dtype=object, shape=(-1,)),\n",
    "            ],\n",
    "            config=ModelConfig(\n",
    "                max_batch_size=64,\n",
    "                batcher=DynamicBatcher(max_queue_delay_microseconds=5000),  # 5ms\n",
    "            ),\n",
    "            strict=True,\n",
    "        )\n",
    "\n",
    "        def _stop_triton(signum, frame):\n",
    "            # The server manager sends SIGTERM to stop the server; this function ensures graceful cleanup.\n",
    "            print(\"SERVER: Received SIGTERM. Stopping Triton server.\")\n",
    "            triton.stop()\n",
    "\n",
    "        signal.signal(signal.SIGTERM, _stop_triton)\n",
    "\n",
    "        print(\"SERVER: Serving inference\")\n",
    "        triton.serve()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "527da1b0",
   "metadata": {},
   "source": [
    "#### Start Triton servers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4142ebfc",
   "metadata": {},
   "source": [
    "The `TritonServerManager` will handle the lifecycle of Triton server instances across the Spark cluster:\n",
    "- Find available ports for HTTP/gRPC/metrics\n",
    "- Deploy a server on each node via stage-level scheduling\n",
    "- Gracefully shutdown servers across nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d522f30",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"ConditionalGeneration\"\n",
    "server_manager = TritonServerManager(model_name=model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c18994c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-07 11:03:44,809 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-07 11:03:44,810 - INFO - Starting 1 servers.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'cb4ae00-lcedt': (2020631, [7000, 7001, 7002])}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Returns {'hostname', (server_pid, [http_port, grpc_port, metrics_port])}\n",
    "server_manager.start_servers(triton_server)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f284eb3",
   "metadata": {},
   "source": [
    "#### Define client function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "237e56dd",
   "metadata": {},
   "source": [
    "Get the hostname -> url mapping from the server manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "826db582",
   "metadata": {},
   "outputs": [],
   "source": [
    "host_to_http_url = server_manager.host_to_http_url  # or server_manager.host_to_grpc_url"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3f58e7b",
   "metadata": {},
   "source": [
    "Define the Triton inference function, which returns a predict function for batch inference through the server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "aff88b3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_fn(model_name, host_to_url):\n",
    "    import socket\n",
    "    import numpy as np\n",
    "    from pytriton.client import ModelClient\n",
    "\n",
    "    url = host_to_url[socket.gethostname()]\n",
    "    print(f\"Connecting to Triton model {model_name} at {url}.\")\n",
    "\n",
    "    def infer_batch(inputs):\n",
    "        with ModelClient(url, model_name, inference_timeout_s=240) as client:\n",
    "            flattened = np.squeeze(inputs).tolist() \n",
    "            # Encode batch\n",
    "            encoded_batch = [[text.encode(\"utf-8\")] for text in flattened]\n",
    "            encoded_batch_np = np.array(encoded_batch, dtype=np.bytes_)\n",
    "            # Run inference\n",
    "            result_data = client.infer_batch(encoded_batch_np)\n",
    "            result_data = np.squeeze(result_data[\"translations\"], -1)\n",
    "            return result_data\n",
    "        \n",
    "    return infer_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "5d10c61c-6102-4d19-8dd6-0c7b5b65343e",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate = predict_batch_udf(partial(triton_fn, model_name=model_name, host_to_url=host_to_http_url),\n",
    "                             return_type=StringType(),\n",
    "                             input_tensor_shapes=[[1]],\n",
    "                             batch_size=32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a85e2ceb",
   "metadata": {},
   "source": [
    "#### Load and preprocess DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "2fa3664e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(text: pd.Series, prefix: str = \"\") -> pd.Series:\n",
    "    @pandas_udf(\"string\")\n",
    "    def _preprocess(text: pd.Series) -> pd.Series:\n",
    "        return pd.Series([prefix + s.split(\".\")[0] for s in text])\n",
    "    return _preprocess(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "5d6c54e7-534d-406f-b8e6-fd592efd0ab2",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.parquet(data_path).limit(256).repartition(8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "dc1bbbe3-4232-49e5-80f6-99976524b73b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:55:37 WARN CacheManager: Asked to cache already cached data.\n"
     ]
    }
   ],
   "source": [
    "input_df = df.select(preprocess(col(\"text\"), \"translate English to French: \").alias(\"input\")).cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e71f07d4",
   "metadata": {},
   "source": [
    "#### Run Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "2e0907da-a5d9-4c3b-9db4-ce5e70ca9bb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 51:==================================================>       (7 + 1) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 10.8 ms, sys: 8.12 ms, total: 18.9 ms\n",
      "Wall time: 30 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "preds = input_df.withColumn(\"preds\", generate(struct(\"input\")))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "9308bdd7-6f67-484d-8b51-dd1e1b2960ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 54:===========================================>              (6 + 2) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.23 ms, sys: 3.43 ms, total: 10.7 ms\n",
      "Wall time: 21.2 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df.withColumn(\"preds\", generate(\"input\"))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "38484ffd-370d-492b-8ca4-9eff9f242a9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 57:===========================================>              (6 + 2) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.81 ms, sys: 12.7 ms, total: 15.5 ms\n",
      "Wall time: 22.3 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = input_df.withColumn(\"preds\", generate(col(\"input\")))\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "ebcb6699-3ac2-4529-ab0f-fab0a5e792da",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 60:>                                                         (0 + 1) / 1]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|                                             input|                                             preds|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "|translate English to French: Doesn't anyone bot...|Ne s'ennuie-t-il pas de vérifier où viennent ce...|\n",
      "|translate English to French: There were two thi...|Il y avait deux choses que j'ai hâte de voir : ...|\n",
      "|translate English to French: I'm rather surpris...|Je suis plutôt surpris que quelqu'un ait trouvé...|\n",
      "|translate English to French: Cultural Vandalism...|Vandalisme culturel La nouvelle production Hall...|\n",
      "|translate English to French: I was at Wrestlema...|J'étais à Wrestlemania VI à Toronto en 10 ans, ...|\n",
      "|translate English to French: This movie has bee...|                       Ce film a été réalisé avant|\n",
      "|translate English to French: [ as a new resolut...|[ en tant que nouvelle résolution pour cette an...|\n",
      "|translate English to French: This movie is over...|Je suis triste de dire que je parviens à regard...|\n",
      "|translate English to French: This show had a pr...|Ce spectacle a eu un début prometteur en l'espè...|\n",
      "|translate English to French: MINOR PLOT SPOILER...|br />br /> Comment ces acteurs talentueux ont-i...|\n",
      "|translate English to French: There is not one c...|Il n'y a pas d'un personnage sur ce sitcom ayan...|\n",
      "|translate English to French: Tommy Lee Jones wa...|Tommy Lee Jones était le meilleur Woodroe et pe...|\n",
      "|translate English to French: My wife rented thi...|Ma femme a loué ce film et n'a jamais pu le voi...|\n",
      "|translate English to French: This is one of tho...|C’est l’une des comédies en étoiles à l’étoile ...|\n",
      "|translate English to French: This excruciatingl...|Ce film excruciant ennuyant et infaillible m’a ...|\n",
      "|translate English to French: you will likely be...|Vous serez probablement très déçu par cette séq...|\n",
      "|translate English to French: If I was British, ...|Si j'étais britannique, je seraitis embarrassé ...|\n",
      "|translate English to French: One of those movie...|Un des films dans lesquels il n'y a pas de gros...|\n",
      "|translate English to French: This show is like ...|Ce spectacle ressemble à l'observation d'une pe...|\n",
      "|                 translate English to French: Sigh|                                             Pesée|\n",
      "+--------------------------------------------------+--------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "preds.show(truncate=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "919e3113-64dd-482a-9233-6607b3f63c1e",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Shut down server on each executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "425d3b28-7705-45ba-8a18-ad34fc895219",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:56:54,506 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-04 13:56:59,695 - INFO - Sucessfully stopped 1 servers.                 \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[True]"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "server_manager.stop_servers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "2dec80ca-7a7c-46a9-97c0-7afb1572f5b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not on_databricks: # on databricks, spark.stop() puts the cluster in a bad state\n",
    "    spark.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f43118ab-fc0a-4f64-a126-4302e615654a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spark-dl-tf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
